import os
import pandas as pd
import numpy as np
from time import time
from tqdm import tqdm
import gc
import warnings

# 仅忽略特定类型的警告
warnings.filterwarnings("ignore", category=UserWarning)

# 设置随机种子，确保实验可复现
np.random.seed(1997)

# 抽象函数：处理点击数据为序列化样本
def process_user_clicks(data, max_len=19, is_training=True):
    user_ids, sequences, targets = [], [], []
    for user_id, records in tqdm(data.groupby('user_id')):
        click_history = records['click_article_id'].values.tolist()
        history_length = len(click_history)

        if history_length <= 1:
            if not is_training:
                user_ids.append(user_id)
                sequences.append(click_history)
            continue

        for i in range(history_length - 1):
            if is_training:
                # 构建训练样本
                if history_length <= max_len + 1 or i < history_length - max_len:
                    user_ids.append(user_id)
                    sequences.append(click_history[:i + 1])
                    targets.append(click_history[i + 1])
            elif i == history_length - 2:
                # 构建验证或测试样本
                user_ids.append(user_id)
                sequences.append(click_history[:i + 1])
                targets.append(click_history[i + 1])

    return user_ids, sequences, targets

# 构建训练和验证样本
def construct_samples_train(train_data, valid_users, max_len=19, save_dir='../samples_tmp/', model_label='srgnn'):
    train_data.sort_values("click_timestamp", inplace=True)
    os.makedirs(save_dir, exist_ok=True)

    # 分割训练集和验证集
    validation_data = train_data[train_data.user_id.isin(valid_users)]
    training_data = train_data[~train_data.user_id.isin(valid_users)]

    # 处理训练样本
    train_user_ids, train_sequences, train_targets = process_user_clicks(training_data, max_len, is_training=True)
    # 处理验证样本
    valid_user_ids, valid_sequences, valid_targets = process_user_clicks(validation_data, max_len, is_training=False)

    # 保存序列化数据
    np.savez(os.path.join(save_dir, f'train_{model_label}.npz'), users=train_user_ids, seqs=train_sequences, targets=train_targets)
    np.savez(os.path.join(save_dir, f'valid_{model_label}.npz'), users=valid_user_ids, seqs=valid_sequences, targets=valid_targets)

    print(f"Train samples: {len(train_user_ids)}, Valid samples: {len(valid_user_ids)}")
    return (train_user_ids, train_sequences, train_targets), (valid_user_ids, valid_sequences, valid_targets)

# 构建训练和测试样本
def construct_samples_test(train, test_data, max_len=19, save_dir='../samples_tmp/'):
    test_data.sort_values("click_timestamp", inplace=True)
    os.makedirs(save_dir, exist_ok=True)

    # 解压训练数据
    train_user_ids, train_sequences, train_targets = list(train[0]), list(train[1]), list(train[2])
    # 处理测试样本
    test_user_ids, test_sequences, _ = process_user_clicks(test_data, max_len, is_training=False)

    # 保存测试样本
    np.savez(os.path.join(save_dir, 'test_samples.npz'), users=test_user_ids, seqs=test_sequences)

    print(f"Train samples: {len(train_user_ids)}, Test samples: {len(test_user_ids)}")
    return (train_user_ids, train_sequences, train_targets), (test_user_ids, test_sequences)

# 主函数示例
if __name__ == "__main__":
    # 配置路径和参数
    train_file_path = '../tcdata/train_click_log.csv'
    valid_user_file_path = 'valid_users.npy'
    save_directory = '../samples_tmp/'

    # 加载数据
    train_click_data = pd.read_csv(train_file_path)
    valid_users = np.load(valid_user_file_path)

    # 构建训练和验证样本
    train, valid = construct_samples_train(train_click_data, valid_users, max_len=19, save_dir=save_directory)

    # 假设加载测试数据
    # test_click_data = pd.read_csv('../tcdata/test_click_log.csv')
    # train, test = construct_samples_test(train, test_click_data, max_len=19, save_dir=save_directory)
